{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/siyuan/miniconda3/envs/search/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/home/siyuan/miniconda3/envs/search/lib/python3.10/site-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils.\n",
      "  warnings.warn(\"Setuptools is replacing distutils.\")\n",
      "/home/siyuan/miniconda3/envs/search/lib/python3.10/site-packages/paddle/jit/dy2static/program_translator.py:712: UserWarning: full_graph=False don't support input_spec arguments. It will not produce any effect.\n",
      "You can set full_graph=True, then you can assign input spec.\n",
      "\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from functools import partial\n",
    "import argparse\n",
    "import os\n",
    "import sys\n",
    "import random\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import hnswlib\n",
    "import paddle\n",
    "import paddle.nn.functional as F\n",
    "from paddlenlp.data import Stack, Tuple, Pad\n",
    "from paddlenlp.datasets import load_dataset, MapDataset\n",
    "from paddlenlp.utils.log import logger\n",
    "import paddlenlp\n",
    "\n",
    "from model import PairwiseMatching\n",
    "from base_model import SemanticIndexBase\n",
    "from data2 import convert_example, create_dataloader\n",
    "from data import convert_pairwise_example\n",
    "from data2 import gen_id2corpus\n",
    "from ann_util import build_index\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下面是准备工作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batchify_fn(samples):\n",
    "    fn = Tuple(\n",
    "        Pad(axis=0, pad_val=tokenizer.pad_token_id, dtype=\"int64\"),  \n",
    "        Pad(axis=0, pad_val=tokenizer.pad_token_type_id, dtype=\"int64\"),  \n",
    "    )\n",
    "\n",
    "    processed_samples = fn(samples) \n",
    "\n",
    "    result = []\n",
    "    for data in processed_samples:\n",
    "        result.append(data) \n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下面是召回的准备工作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_file = \"corpus.csv\" \n",
    "\n",
    "id2corpus = gen_id2corpus(corpus_file) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Cannot open file",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m output_emb_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m256\u001b[39m\n\u001b[1;32m      3\u001b[0m final_index \u001b[38;5;241m=\u001b[39m hnswlib\u001b[38;5;241m.\u001b[39mIndex(space\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mip\u001b[39m\u001b[38;5;124m\"\u001b[39m, dim\u001b[38;5;241m=\u001b[39moutput_emb_size) \n\u001b[0;32m----> 4\u001b[0m \u001b[43mfinal_index\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_index\u001b[49m\u001b[43m(\u001b[49m\u001b[43msave_index_path\u001b[49m\u001b[43m)\u001b[49m \n",
      "\u001b[0;31mRuntimeError\u001b[0m: Cannot open file"
     ]
    }
   ],
   "source": [
    "save_index_path = \"final_index.bin\" \n",
    "output_emb_size = 256\n",
    "final_index = hnswlib.Index(space=\"ip\", dim=output_emb_size) \n",
    "final_index.load_index(save_index_path) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[2024-01-29 11:10:59,336] [    INFO]\u001b[0m - We are using (<class 'paddlenlp.transformers.ernie.tokenizer.ErnieTokenizer'>, False) to load 'ernie-1.0'.\u001b[0m\n",
      "\u001b[32m[2024-01-29 11:10:59,339] [    INFO]\u001b[0m - Already cached /home/siyuan/.paddlenlp/models/ernie-1.0/vocab.txt\u001b[0m\n",
      "\u001b[32m[2024-01-29 11:10:59,355] [    INFO]\u001b[0m - tokenizer config file saved in /home/siyuan/.paddlenlp/models/ernie-1.0/tokenizer_config.json\u001b[0m\n",
      "\u001b[32m[2024-01-29 11:10:59,356] [    INFO]\u001b[0m - Special tokens file saved in /home/siyuan/.paddlenlp/models/ernie-1.0/special_tokens_map.json\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "model_name = \"ernie-1.0\"\n",
    "\n",
    "tokenizer = paddlenlp.transformers.AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans_func = partial(convert_example, tokenizer=tokenizer, max_seq_length=60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_model = paddlenlp.transformers.AutoModel.from_pretrained(model_name)\n",
    "\n",
    "model = SemanticIndexBase(pretrained_model, output_emb_size=256)\n",
    "\n",
    "params_path = \"recall_model_state.pdparams\" \n",
    "\n",
    "if params_path and os.path.isfile(params_path): \n",
    "    state_dict = paddle.load(params_path) \n",
    "    model.set_dict(state_dict) \n",
    "    print(\"Loaded parameters from %s\" % params_path) \n",
    "else:\n",
    "    raise ValueError(\"Please set params_path with correct pretrained model file\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#下面是排序的准备工作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@paddle.no_grad()\n",
    "def predict(model, data_loader):\n",
    "    model.eval()\n",
    "\n",
    "    batch_probs = []\n",
    "    for batch_data in data_loader:\n",
    "        input_ids, token_type_ids = batch_data\n",
    "        batch_prob = model.predict(input_ids=input_ids, token_type_ids=token_type_ids).numpy()\n",
    "        batch_probs.append(batch_prob)\n",
    "\n",
    "    conca_batch_probs = np.concatenate(batch_probs, axis=0) \n",
    "\n",
    "    return conca_batch_probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_model_name = \"ernie-3.0-medium-zh\"\n",
    "\n",
    "rank_pretrained_model = paddlenlp.transformers.AutoModel.from_pretrained(rank_model_name)\n",
    "\n",
    "rank_model = PairwiseMatching(rank_pretrained_model)\n",
    "\n",
    "rank_params_path = \"rank_model_state.pdparams\" \n",
    "\n",
    "if rank_params_path and os.path.isfile(rank_params_path): \n",
    "    state_dict = paddle.load(rank_params_path)\n",
    "    rank_model.set_dict(state_dict)\n",
    "    print(\"Loaded parameters from %s\" % params_path)\n",
    "else:\n",
    "    raise ValueError(\"Please set --params_path with correct pretrained model file\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_tokenizer = paddlenlp.transformers.AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_trans_func = partial(convert_pairwise_example, tokenizer=rank_tokenizer, max_seq_length=64, phase=\"predict\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#下面开始进行搜索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_query = input(\"请输入需要查询的文献\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "querys = []\n",
    "querys.append({\"text\": input_query})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(querys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_ds = MapDataset(querys)\n",
    "\n",
    "query_batch_sampler = paddle.io.BatchSampler(query_ds, batch_size=1, shuffle=False)\n",
    "\n",
    "query_data_loader = paddle.io.DataLoader(dataset=query_ds.map(trans_func), batch_sampler=query_batch_sampler, collate_fn=batchify_fn, return_list=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_embedding = model.get_semantic_embedding(query_data_loader) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recall_data = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下面正式开始召回\n",
    "\n",
    "for batch_index, batch_query_embedding in enumerate(query_embedding): \n",
    "    recalled_idx, cosine_sims = final_index.knn_query(batch_query_embedding.numpy(), 50)                                                                           \n",
    "\n",
    "    for idx, doc_idx in enumerate(recalled_idx[0]): \n",
    "        recall_data.append({\"query\": input_query, \"title\": id2corpus[doc_idx]})\n",
    "        print( \"{}\\t{}\\n\".format(id2corpus[doc_idx], 1.0 - cosine_sims[0][idx] )                                                                         \n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in recall_data:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_ds = MapDataset(recall_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_ds_copy = copy.deepcopy(rank_ds)\n",
    "\n",
    "rank_batch_sampler = paddle.io.BatchSampler(rank_ds, batch_size=16, shuffle=False)\n",
    "\n",
    "rank_data_loader = paddle.io.DataLoader(dataset=rank_ds.map(rank_trans_func), batch_sampler=rank_batch_sampler, collate_fn=batchify_fn, return_list=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_probs = predict(rank_model, rank_data_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_result = [] \n",
    "\n",
    "for idx, prob in enumerate(y_probs): \n",
    "    text_pair = rank_ds_copy[idx] \n",
    "    text_pair[\"pred_prob\"] = prob[0] \n",
    "    rank_result.append(text_pair) \n",
    "    print(text_pair) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_result.sort(key=lambda x: x['pred_prob'], reverse=True)\n",
    "\n",
    "for i in rank_result:\n",
    "    print(\"{}\\t{}\".format(i['title'], i['pred_prob']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下面是可以整体运行的代码\n",
    "\n",
    "input_query = input(\"请输入需要查询的文献\")\n",
    "\n",
    "querys = []\n",
    "querys.append({\"text\": input_query})\n",
    "\n",
    "query_ds = MapDataset(querys)\n",
    "\n",
    "query_data_loader = create_dataloader(\n",
    "    query_ds, mode=\"predict\", batch_size=1, batchify_fn=batchify_fn, trans_fn=trans_func\n",
    ")\n",
    "\n",
    "query_embedding = model.get_semantic_embedding(query_data_loader) \n",
    "\n",
    "recall_data = []\n",
    "\n",
    "# 下面正式开始召回\n",
    "\n",
    "for batch_index, batch_query_embedding in enumerate(query_embedding): \n",
    "    recalled_idx, cosine_sims = final_index.knn_query(batch_query_embedding.numpy(), 50) \n",
    "    for idx, doc_idx in enumerate(recalled_idx[0]): \n",
    "        recall_data.append({\"query\": input_query, \"title\": id2corpus[doc_idx], \"sim\": 1.0 - cosine_sims[0][idx]})\n",
    "        print( \"{}\\t{}\\n\".format(id2corpus[doc_idx], 1.0 - cosine_sims[0][idx] ) \n",
    "        )\n",
    "\n",
    "rank_ds = MapDataset(recall_data)\n",
    "\n",
    "\n",
    "rank_ds_copy = copy.deepcopy(rank_ds)\n",
    "\n",
    "rank_data_loader = create_dataloader(\n",
    "    rank_ds_copy, mode=\"predict\", batch_size=16, batchify_fn=batchify_fn, trans_fn=rank_trans_func\n",
    ")\n",
    "\n",
    "y_probs = predict(rank_model, rank_data_loader)\n",
    "\n",
    "rank_result = []\n",
    "\n",
    "for idx, prob in enumerate(y_probs):\n",
    "    text_pair = rank_ds[idx]\n",
    "    text_pair[\"pred_prob\"] = prob[0]\n",
    "    rank_result.append(text_pair)\n",
    "\n",
    "rank_result.sort(key=lambda x: x['pred_prob'], reverse=True)\n",
    "\n",
    "for i in rank_result:\n",
    "    print(\"{}\\t{}\".format(i['title'], i['pred_prob']))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "search",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
